package main

import (
	"flag"
	"fmt"
	"net"
	"os"

	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/gorilla/websocket"
	"github.com/op/go-logging"
	"io"
	"net/http"
	"path"
	"strings"
	"sync"
	"time"
)

const ModuleName = "rproxy"

var (
	log          = logging.MustGetLogger(ModuleName)
	daemon       = flag.Bool("d", false, "run as a daemon")
	bindHost     = flag.String("bind", "0.0.0.0:1", "the interface for binding")
	upstreamHost = flag.String("upstream", "0.0.0.0:1", "the upstream host for relaying")
	localHost    = flag.String("local", "0.0.0.0", "the local host for relaying")
	debug        = flag.Bool("debug", false, "run in debug mode")
	logPath      = flag.String("log", "", "the log path")
	wsServer     = flag.Bool("ws", false, "run as a websocket server")
)

type proxyInfo struct {
	conn         net.Conn
	upstreamAddr *net.TCPAddr
	localAddr    *net.TCPAddr
}

type bridgeInfo struct {
	clientConn   net.Conn
	upstreamConn net.Conn
}

const HEADSIZE = 2

func syncRead(info bridgeInfo, closeEvt chan<- string) {
	defer func() {
		if err := recover(); err != nil {
			log.Error("caught an error:", err)
			closeEvt <- fmt.Sprintf("failed with error:%s", err)
		}
	}()

	var reason string

	originAddr := info.clientConn.RemoteAddr().(*net.TCPAddr)
	originIP := originAddr.IP.To4()
	packedIP := uint32(originIP[0])<<24 | uint32(originIP[1])<<16 |
		uint32(originIP[2])<<8 | uint32(originIP[3])

	head := make([]byte, HEADSIZE)
	headReader := bytes.NewReader(head)

	clientConn := bufio.NewReader(info.clientConn)
	upstreamConn := bufio.NewWriter(info.upstreamConn)
	for {
		_, err := io.ReadAtLeast(clientConn, head, HEADSIZE)
		if err != nil {
			reason = fmt.Sprintf("failed to read data from client (%s)", err)
			break
		}

		headReader.Reset(head)

		packetSize := int16(0)
		err = binary.Read(headReader, binary.BigEndian, &packetSize)
		if err != nil {
			reason = fmt.Sprintf("failed to read head from client (%s)", err)
			break
		}

		log.Debugf("received from[%d.%d.%d.%d:%d] packet length:%d",
			originIP[0], originIP[1], originIP[2], originIP[3], packedIP, packetSize)

		//增加包的大小,把ip塞进去
		err = binary.Write(upstreamConn, binary.BigEndian, packetSize+4)
		if err != nil {
			reason = fmt.Sprintf("failed to write head (%s)", err)
			break
		}
		binary.Write(upstreamConn, binary.BigEndian, packedIP)

		if packetSize > 0 {
			n, err := io.CopyN(upstreamConn, clientConn, int64(packetSize))
			if err != nil {
				reason = fmt.Sprintf("failed to copy (%d) (%d) (%s)", packetSize, n, err)
				break
			}
		}

		upstreamConn.Flush()
	}
	closeEvt <- reason
}

func syncWrite(info bridgeInfo, closeEvt chan<- string) {
	defer func() {
		if err := recover(); err != nil {
			log.Error("caught an error:", err)
			closeEvt <- fmt.Sprintf("failed with error:%s", err)
		}
	}()

	var reason string

	for {
	    len, err := io.Copy(info.clientConn, info.upstreamConn)
		if err != nil {
		    reason = fmt.Sprintf("failed to read data from upstream (%s)", err)
			break
		}
		if len == 0 {
			reason = fmt.Sprintf("connection closed!")
			break
		}
		log.Debugf("relay message successfully, len(%d)", len)
    }

	closeEvt <- reason
}

func proxyToUpstream(info *proxyInfo) {
	defer func() {
		if err := recover(); err != nil {
			log.Error("caught an error:", err)
		}
	}()

	log.Infof("accept the new connection [%s]", info.conn.RemoteAddr().String())

	conn, err := net.DialTCP("tcp", info.localAddr, info.upstreamAddr)
	if err != nil {
		log.Errorf("failed to connect upstream [%s]->[%s]",
			info.conn.RemoteAddr().String(), info.upstreamAddr.String())
		return
	}
	conn.SetKeepAlive(true)
	conn.SetKeepAlivePeriod(3 * time.Minute)

	readCloseEvt := make(chan string)
	writeCloseEvt := make(chan string)

	bInfo := bridgeInfo{
		clientConn:   info.conn,
		upstreamConn: conn,
	}

	go syncRead(bInfo, readCloseEvt)
	go syncWrite(bInfo, writeCloseEvt)

	var reason string
	var otherChan chan string
	closeFunc := func() {
		bInfo.clientConn.Close()
		bInfo.upstreamConn.Close()
	}
	select {
	case reason = <-readCloseEvt:
		otherChan = writeCloseEvt
		closeFunc()
	case reason = <-writeCloseEvt:
		otherChan = readCloseEvt
		closeFunc()
	}

	<-otherChan
	log.Infof("disconnect a connect([%s]) successfully! reason(%s)", bInfo.clientConn.RemoteAddr().String(), reason)
}

func StartServer(bindAddr, localAddr, upstreamAddr string) error {
	l, err := net.Listen("tcp4", bindAddr)
	if err != nil {
		return errors.New(fmt.Sprintf("failed to listen server (%s), error(%s)", bindAddr, err))
	}

	defer l.Close()

	log.Infof("now listening on (%s)", bindAddr)

	upAddr, err := net.ResolveTCPAddr("tcp", upstreamAddr)
	if err != nil {
		return errors.New(fmt.Sprintf("failed to resolve upstream address (%s),error(%s)", upstreamAddr, err))
	}

	lAddr, err := net.ResolveTCPAddr("tcp", localAddr)
	if err != nil {
		return errors.New(fmt.Sprintf("failed to resolve local address (%s),error(%s)", localAddr, err))
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			log.Errorf("failed to accept the new connection: %s", err)
		}

		go proxyToUpstream(&proxyInfo{
			conn:         conn,
			upstreamAddr: upAddr,
			localAddr:    lAddr,
		})
	}

	return nil
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 4096,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type WebsocketHandler struct {
	upstreamAddr *net.TCPAddr
	localAddr    *net.TCPAddr
}

func (ws *WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	clientConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error(err)
		return
	}

	clientIPStr := strings.TrimSpace(r.Header.Get("Ali-CDN-Real-IP"))
	if clientIPStr == "" {
		params := strings.Split(r.Header.Get("X-Forwarded-For"), ",")
		if len(params) > 0 {
			clientIPStr = strings.TrimSpace(params[0])
		}
	}

	if clientIPStr == "" {
		clientIPStr = "0.0.0.0"
	}

	log.Debugf("accept new connection,real ip is (%s),from proxy (%s)", clientIPStr, clientConn.RemoteAddr().String())
	upstreamConn, err := net.DialTCP("tcp", ws.localAddr, ws.upstreamAddr)
	if err != nil {
		log.Errorf("failed to connect upstream [%s]", ws.upstreamAddr.String())
		return
	}
	upstreamConn.SetKeepAlive(true)
	upstreamConn.SetKeepAlivePeriod(3 * time.Minute)

	packedIP := uint32(0)
	tmpAddr, err := net.ResolveIPAddr("ip", clientIPStr)

	if err != nil {
		log.Error(clientIPStr, err)
	} else {
		originIP := tmpAddr.IP.To4()
		packedIP = uint32(originIP[0])<<24 | uint32(originIP[1])<<16 |
			uint32(originIP[2])<<8 | uint32(originIP[3])
	}

	monitor := sync.WaitGroup{}
	monitor.Add(2)

	go func() {
		defer func() {
			if err := recover(); err != nil {
				log.Error("caught an error:", err)
			}
			upstreamConn.Close()
			monitor.Done()
		}()

		upstreamWriter := bufio.NewWriter(upstreamConn)

		for {
			messageType, p, err := clientConn.ReadMessage()
			if err != nil {
				log.Error(err)
				return
			}

			clientReader := bytes.NewReader(p)

			packetSize := int16(0)
			err = binary.Read(clientReader, binary.BigEndian, &packetSize)
			if err != nil {
				log.Errorf("failed to read head from client (%s)", err)
				return
			}

			log.Debugf("received from packedIP(%d) message type<%d>,data length <%d>, packet length:%d",
				packedIP, messageType, len(p), packetSize)

			//增加包的大小,把ip塞进去
			err = binary.Write(upstreamWriter, binary.BigEndian, packetSize+4)
			if err != nil {
				log.Errorf("failed to write head (%s)", err)
				return
			}
			binary.Write(upstreamWriter, binary.BigEndian, packedIP)

			if packetSize > 0 {
				n, err := io.CopyN(upstreamWriter, clientReader, int64(packetSize))
				if err != nil {
					log.Errorf("failed to copy (%d) (%d)", packetSize, n)
					return
				}
			}

			upstreamWriter.Flush()
		}
	}()

	go func() {
		defer func() {
			if err := recover(); err != nil {
				log.Error("caught an error:", err)
			}
			clientConn.Close()
			monitor.Done()
		}()

		head := make([]byte, HEADSIZE)
		headReader := bytes.NewReader(head)

		upstreamReader := bufio.NewReader(upstreamConn)

		for {
			n, err := io.ReadAtLeast(upstreamReader, head, len(head))
			if err != nil {
				log.Error(err)
				return
			}

			if n != len(head) {
				log.Errorf("invalid head %d", n)
				return
			}

			headReader.Reset(head)

			packetSize := int16(0)
			err = binary.Read(headReader, binary.BigEndian, &packetSize)
			if err != nil {
				log.Error(err)
				return
			}

			body := make([]byte, packetSize+HEADSIZE)
			body[0] = head[0]
			body[1] = head[1]
			n, err = io.ReadAtLeast(upstreamReader, body[HEADSIZE:], int(packetSize))
			if err != nil {
				log.Error(err)
				return
			}

			if n != int(packetSize) {
				log.Error("invalid body ", n, packetSize)
				return
			}

			if err := clientConn.WriteMessage(websocket.BinaryMessage, body); err != nil {
				log.Error(err)
				return
			}
		}
	}()

	monitor.Wait()

	log.Debug("connection completed!")
}

func StartWebSocketServer(bindAddr, localAddr, upstreamAddr string) error {
	var connType = "tcp4"
	if path.IsAbs(bindAddr) {
		connType = "unix"
	}

	l, err := net.Listen(connType, bindAddr)
	if err != nil {
		return errors.New(fmt.Sprintf("failed to listen server (%s), error(%s)", bindAddr, err))
	}

	defer l.Close()

	if path.IsAbs(bindAddr) {
		if err := os.Chmod(bindAddr, 0777); err != nil {
			log.Error(err)
			return err
		}
	}

	log.Infof("now listening web connection on (%s)", bindAddr)

	upAddr, err := net.ResolveTCPAddr("tcp", upstreamAddr)
	if err != nil {
		return errors.New(fmt.Sprintf("failed to resolve upstream address (%s),error(%s)", upstreamAddr, err))
	}

	lAddr, err := net.ResolveTCPAddr("tcp", localAddr)
	if err != nil {
		return errors.New(fmt.Sprintf("failed to resolve local address (%s),error(%s)", localAddr, err))
	}

	server := http.Server{
		Handler: &WebsocketHandler{
			upstreamAddr: upAddr,
			localAddr:    lAddr,
		},
	}

	if err := server.Serve(l); err != nil {
		log.Error(err)
	}

	return nil
}

func setupLogger() error {
	loggerWriter := os.Stdout
	if *logPath != "" {
		var err error
		loggerWriter, err = os.OpenFile(*logPath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0644)
		if err != nil {
			return err
		}
	}

	format := logging.MustStringFormatter(
		`%{color}[%{time:15:04:05.000}] [%{level:.4s}] %{shortfile} [%{shortfunc}] %{color:reset} %{message}`,
	)

	logLevel := logging.INFO
	if *debug {
		logLevel = logging.DEBUG
	}

	logging.SetFormatter(format)
	backend := logging.NewLogBackend(loggerWriter, "", 0)
	backendLeveled := logging.AddModuleLevel(backend)
	backendLeveled.SetLevel(logLevel, ModuleName)

	logging.SetBackend(backendLeveled)

	return nil
}

func main() {
	flag.Parse()

	if err := setupLogger(); err != nil {
		panic(err)
	}

	bindAddr := *bindHost
	upstreamAddr := *upstreamHost
	localAddr := *localHost

	if *wsServer {
		log.Info("run as a websocket server...")
		if err := StartWebSocketServer(bindAddr, localAddr, upstreamAddr); err != nil {
			log.Info("failed to run server", bindAddr, upstreamAddr, err)
			os.Exit(-1)
		}
	} else {
		log.Infof("run as a tcp proxy server...")
		if err := StartServer(bindAddr, localAddr, upstreamAddr); err != nil {
			log.Info("failed to run server", bindAddr, upstreamAddr, err)
			os.Exit(-1)
		}
	}

	log.Info("now server exit")
}
